package jxmap

import (
	"encoding/binary"
	"hash/fnv"
	"math"
	"sync"

	"github.com/vmihailenco/msgpack/v5"
)

// JxMap 高性能并发字典
type JxMap[K comparable, V any] struct {
	shardCount int // 新增分片数量字段
	shards     []*shard[K, V]
}

type shard[K comparable, V any] struct {
	sync.RWMutex
	data map[K]V
}

// MarshalMsgpack 实现 msgpack 序列化
func (m *JxMap[K, V]) MarshalMsgpack() ([]byte, error) {
	// 创建一个临时 map 用于存储所有分片数据
	temp := make(map[K]V)

	// 遍历所有分片，合并数据
	m.Range(func(key K, value V) bool {
		temp[key] = value
		return true
	})

	// 使用 msgpack 序列化临时 map
	return msgpack.Marshal(temp)
}

// UnmarshalMsgpack 实现 msgpack 反序列化
func (m *JxMap[K, V]) UnmarshalMsgpack(b []byte) error {
	// 先反序列化到一个临时 map
	var temp map[K]V
	if err := msgpack.Unmarshal(b, &temp); err != nil {
		return err
	}

	// 清空当前 map
	for _, sh := range m.shards {
		sh.Lock()
		sh.data = make(map[K]V)
		sh.Unlock()
	}

	// 将数据填充到分片中
	for k, v := range temp {
		shard := m.getShard(k)
		shard.Lock()
		shard.data[k] = v
		shard.Unlock()
	}

	return nil
}

// New 创建新的并发字典（默认32个分片）
func New[K comparable, V any]() *JxMap[K, V] {
	return NewWithShards[K, V](32)
}

// NewWithShards 指定分片数量创建
func NewWithShards[K comparable, V any](shardCount int) *JxMap[K, V] {
	m := &JxMap[K, V]{
		shardCount: shardCount, // 初始化分片数量
		shards:     make([]*shard[K, V], shardCount),
	}
	for i := range m.shards {
		m.shards[i] = &shard[K, V]{
			data: make(map[K]V),
		}
	}
	return m
}

// 获取分片（优化版）
func (m *JxMap[K, V]) getShard(key K) *shard[K, V] {
	hash := hashKey(key)
	return m.shards[hash%uint32(m.shardCount)]
}
func hashKey[K comparable](key K) uint32 {
	h := fnv.New32a()
	switch v := any(key).(type) {
	case string:
		h.Write([]byte(v))
	case int:
		var buf [8]byte
		binary.BigEndian.PutUint64(buf[:], uint64(v))
		h.Write(buf[:])
	case int8:
		var buf [1]byte
		buf[0] = byte(v)
		h.Write(buf[:])
	case int16:
		var buf [2]byte
		binary.BigEndian.PutUint16(buf[:], uint16(v))
		h.Write(buf[:])
	case int32:
		var buf [4]byte
		binary.BigEndian.PutUint32(buf[:], uint32(v))
		h.Write(buf[:])
	case int64:
		var buf [8]byte
		binary.BigEndian.PutUint64(buf[:], uint64(v))
		h.Write(buf[:])
	case uint:
		var buf [8]byte
		binary.BigEndian.PutUint64(buf[:], uint64(v))
		h.Write(buf[:])
	case uint8:
		var buf [1]byte
		buf[0] = byte(v)
		h.Write(buf[:])
	case uint16:
		var buf [2]byte
		binary.BigEndian.PutUint16(buf[:], uint16(v))
		h.Write(buf[:])
	case uint32:
		var buf [4]byte
		binary.BigEndian.PutUint32(buf[:], uint32(v))
		h.Write(buf[:])
	case uint64:
		var buf [8]byte
		binary.BigEndian.PutUint64(buf[:], v)
		h.Write(buf[:])
	case float32:
		var buf [4]byte
		binary.BigEndian.PutUint32(buf[:], math.Float32bits(float32(v)))
		h.Write(buf[:])
	case float64:
		var buf [8]byte
		binary.BigEndian.PutUint64(buf[:], math.Float64bits(float64(v)))
		h.Write(buf[:])
	case bool:
		if v {
			h.Write([]byte{1})
		} else {
			h.Write([]byte{0})
		}
	default:
		panic("unsupported key type")
	}
	return h.Sum32()
}

// Set 添加/更新键值对
func (m *JxMap[K, V]) Set(key K, value V) {
	shard := m.getShard(key)
	shard.Lock()
	defer shard.Unlock()
	shard.data[key] = value
}

// Get 获取键值
func (m *JxMap[K, V]) Get(key K) (V, bool) {
	shard := m.getShard(key)
	shard.RLock()
	defer shard.RUnlock()
	val, ok := shard.data[key]
	return val, ok
}

// Delete 删除键
func (m *JxMap[K, V]) Delete(key K) {
	shard := m.getShard(key)
	shard.Lock()
	defer shard.Unlock()
	delete(shard.data, key)
}

// Size 获取元素总数
func (m *JxMap[K, V]) Size() int {
	var count int
	var wg sync.WaitGroup
	wg.Add(len(m.shards))

	for _, s := range m.shards {
		go func(sh *shard[K, V]) {
			defer wg.Done()
			sh.RLock()
			count += len(sh.data)
			sh.RUnlock()
		}(s)
	}
	wg.Wait()
	return count
}
func (m *JxMap[K, V]) Range(f func(key K, value V) bool) {
	for _, sh := range m.shards {
		sh.RLock()
		for k, v := range sh.data {
			if !f(k, v) {
				sh.RUnlock()
				return
			}
		}
		sh.RUnlock()
	}
}

// Update 更新键值对
func (m *JxMap[K, V]) Update(key K, updateFunc func(V) V) bool {
	shard := m.getShard(key)
	shard.Lock()
	defer shard.Unlock()

	if val, ok := shard.data[key]; ok {
		newVal := updateFunc(val)
		shard.data[key] = newVal
		return true
	}
	return false
}
